
import numpy as np


def head_weights_prune_per_row(np_data, keep_weights=True, prune_ratio=0.1):

    dim_0, dim_1 = np_data.shape[0], np_data.shape[1]

    num_to_prune = min(int(dim_1 * prune_ratio), dim_1 - 1)
    print("num_to_prune: ", num_to_prune)

    # head weights 每行将分数排名最后的几个记为 -inf
    head_inds_to_prune = np.argsort(np_data)[:, 0:num_to_prune:1]
    head_inds_to_keep = np.argsort(np_data)[:, -1:- (dim_1 - num_to_prune + 1):-1]

    for i in range(dim_0):
        np_data[i, head_inds_to_prune[i]] = - np.inf

    if not keep_weights:
        for i in range(dim_0):
            np_data[i, head_inds_to_keep[i]] = np.inf

    return np_data


def head_weights_prune_uniform(np_data, keep_weights=True, prune_ratio=0.1):
    pass





if __name__ == "__main__":
    np_data = np.random.rand(3, 3)
    print(np_data)
    np_data = head_weights_prune(np_data, keep_weights=True, prune_ratio=0.4)
    print(np_data)

    np_data = np.random.rand(3, 3)
    print(np_data)
    np_data = head_weights_prune(np_data, keep_weights=False, prune_ratio=0.4)
    print(np_data)








